#!/usr/bin/env python
# encoding: utf-8
'''
@author: songyunlong
@contact: 1243049371@qq.com
@software: Pycharm
@file: Training
@time: 2019/8/4 下午9:34
'''
from Gannet import Generator, Discriminator, DiscriminatorGan
import tensorflow as tf
import os
import numpy as np

class TrainingModel:
    @staticmethod
    def dataoutput(iteration, batch_size, *dataset):
        ''''''
        for i in range(iteration):
            start = 0
            while start != dataset[0].shape[0]:
                features_batch = dataset[0][start:start+batch_size, :]
                labels_batch = dataset[-1][start:start+batch_size, :]
                yield features_batch, labels_batch
                start += batch_size

    def __init__(self, realdataset, realdataset_test):
        ''''''
        self.__x = realdataset
        self.__xtest = realdataset_test
        self.__discriminator_gan = DiscriminatorGan()
        self.__generator = Generator()
        self.__discriminator = Discriminator()

    def gan_disc(self):
        ''''''
        gd_optimizer = tf.keras.optimizers.RMSprop(lr=1e-2, clipvalue=1.0, decay=1e-8)
        self.__discriminator_gan.compile(optimizer=gd_optimizer, loss='binary_crossentropy', metrics=['acc'])

    def disc(self):
        ''''''
        d_optimizer = tf.keras.optimizers.RMSprop(lr=1e-2, clipvalue=1.0, decay=1e-8)
        self.__discriminator.compile(optimizer=d_optimizer, loss='binary_crossentropy', metrics=['acc'])

    def training(self):
        ''''''
        ITERATION = 10000
        BATCH_SIZE = 500
        RANDOM_LEN = 32
        # 设置训练轮数标志
        iter_num = 0
        #编译模型
        self.disc()
        self.gan_disc()
        # 迭代训练gan网络
        for realfeature_batch, _ in TrainingModel.dataoutput(
                ITERATION, BATCH_SIZE, self.__x[:, :-1], self.__x[:, -1][:, np.newaxis]):
            # 获得正态分布采样的噪声向量
            random_vector_train = np.random.normal(loc=0, scale=0.1, size=(BATCH_SIZE, RANDOM_LEN))
            random_vector_test = np.random.normal(loc=0, scale=0.1, size=(self.__xtest.shape[0], RANDOM_LEN))
            # 将噪声向量通过生成器映射为生成数据
            gen_data = self.__generator.predict(random_vector_train)
            # print(gen_data.shape, realfeature_batch.shape)
            batch_data = np.vstack((realfeature_batch, gen_data))
            batch_label = np.vstack((np.ones(shape=(realfeature_batch.shape[0], 1), dtype=np.float32),
                                     np.zeros(shape=(gen_data.shape[0], 1), dtype=np.float32)))
            # 在标签上添加随机噪声
            batch_label += 0.05 * np.random.random(size=batch_label.shape)
            # 训练判别器
            d_loss, _ = self.__discriminator.train_on_batch(x=batch_data, y=batch_label)

            # 重新获得正态分布采样的噪声向量
            random_vector_train2 = np.random.normal(loc=0, scale=0.1, size=(BATCH_SIZE, RANDOM_LEN))
            # random_vector_test2 = np.random.normal(loc=0, scale=0.1, size=(self.__xtest.shape[0], RANDOM_LEN))
            # 创建假数据标签
            batch_label2 = np.ones(shape=(random_vector_train2.shape[0], 1), dtype=np.float32)
            dg_loss, _ = self.__discriminator_gan.train_on_batch(x=random_vector_train2, y=batch_label2)
            if not iter_num % (100): #100*BATCH_SIZE
                #生成真假数据混合判别测试特征矩阵
                feature_gen = self.__generator.predict(random_vector_test)
                # print(feature_gen.shape, self.__xtest.shape)
                test_feature = np.vstack((self.__xtest, feature_gen))
                test_labels = np.vstack((np.ones(shape=(self.__xtest.shape[0], 1), dtype=np.float32),
                                         np.zeros(shape=(feature_gen.shape[0], 1), dtype=np.float32)))
                _, acc1 = self.__discriminator.evaluate(x=test_feature, y=test_labels)
                _, acc2 = self.__discriminator_gan.evaluate(x=random_vector_test,
                                                            y=np.ones(shape=(random_vector_test.shape[0], 1), dtype=np.float32))
            #     print('第%s轮d_loss: %s, dg_loss: %s, d_acc: %s, dg_acc: %s' % (iter_num, d_loss, dg_loss, acc1, acc2))
            iter_num += 1


if __name__ == '__main__':
    pass
